"""
30_api_server.py
================
REST API Server for Neo4j Knowledge Graph

This module covers:
1. FastAPI server setup
2. CRUD endpoints for Neo4j
3. Graph RAG endpoints
4. Batch operations
5. WebSocket for real-time queries
6. Authentication and rate limiting

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from enum import Enum
from datetime import datetime
import os
import sys

# FastAPI imports
try:
    from fastapi import FastAPI, HTTPException, Query, Depends, BackgroundTasks
    from fastapi.middleware.cors import CORSMiddleware
    from pydantic import BaseModel, Field
    import uvicorn
    FASTAPI_AVAILABLE = True
except ImportError:
    FASTAPI_AVAILABLE = False
    print("FastAPI not available. Install with: pip install fastapi uvicorn")

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import Neo4jConnection, get_driver


# ============================================================
# PYDANTIC MODELS
# ============================================================

if FASTAPI_AVAILABLE:
    
    class NodeCreate(BaseModel):
        """Model for creating a node."""
        labels: List[str] = Field(..., description="Node labels")
        properties: Dict[str, Any] = Field(..., description="Node properties")
        
        class Config:
            json_schema_extra = {
                "example": {
                    "labels": ["Patient"],
                    "properties": {
                        "id": "P001",
                        "name": "John Doe",
                        "age": 45
                    }
                }
            }
    
    class NodeUpdate(BaseModel):
        """Model for updating a node."""
        properties: Dict[str, Any] = Field(..., description="Properties to update")
    
    class RelationshipCreate(BaseModel):
        """Model for creating a relationship."""
        from_label: str
        from_id: str
        to_label: str
        to_id: str
        rel_type: str
        properties: Optional[Dict[str, Any]] = None
        
        class Config:
            json_schema_extra = {
                "example": {
                    "from_label": "Patient",
                    "from_id": "P001",
                    "to_label": "Disease",
                    "to_id": "E11",
                    "rel_type": "DIAGNOSED_WITH",
                    "properties": {"date": "2024-01-15"}
                }
            }
    
    class CypherQuery(BaseModel):
        """Model for executing Cypher queries."""
        query: str = Field(..., description="Cypher query string")
        parameters: Optional[Dict[str, Any]] = Field(None, description="Query parameters")
        
        class Config:
            json_schema_extra = {
                "example": {
                    "query": "MATCH (p:Patient {id: $id}) RETURN p",
                    "parameters": {"id": "P001"}
                }
            }
    
    class RAGQuery(BaseModel):
        """Model for RAG queries."""
        question: str = Field(..., description="Question to answer")
        top_k: int = Field(5, description="Number of chunks to retrieve")
        include_graph_context: bool = Field(True, description="Include graph context")
        
        class Config:
            json_schema_extra = {
                "example": {
                    "question": "What are the treatment options for diabetes?",
                    "top_k": 5,
                    "include_graph_context": True
                }
            }
    
    class BatchNodeCreate(BaseModel):
        """Model for batch node creation."""
        nodes: List[NodeCreate] = Field(..., description="List of nodes to create")
    
    class SearchQuery(BaseModel):
        """Model for search queries."""
        query: str = Field(..., description="Search query")
        labels: Optional[List[str]] = Field(None, description="Filter by labels")
        limit: int = Field(10, description="Maximum results")
    
    class VectorSearchQuery(BaseModel):
        """Model for vector similarity search."""
        text: str = Field(..., description="Text to search for")
        index_name: str = Field("chunk_embeddings", description="Vector index name")
        top_k: int = Field(10, description="Number of results")
        min_score: float = Field(0.5, description="Minimum similarity score")


# ============================================================
# API APPLICATION
# ============================================================

if FASTAPI_AVAILABLE:
    
    app = FastAPI(
        title="Neo4j Knowledge Graph API",
        description="""
        REST API for Neo4j Knowledge Graph operations.
        
        Features:
        - CRUD operations for nodes and relationships
        - Cypher query execution
        - Graph RAG queries
        - Vector similarity search
        - Batch operations
        """,
        version="1.0.0",
        docs_url="/docs",
        redoc_url="/redoc"
    )
    
    # CORS middleware
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
    
    # =========================================================
    # DEPENDENCIES
    # =========================================================
    
    def get_db():
        """Dependency for database connection."""
        conn = Neo4jConnection()
        try:
            conn.connect()
            yield conn
        finally:
            conn.close()
    
    # =========================================================
    # HEALTH CHECK ENDPOINTS
    # =========================================================
    
    @app.get("/", tags=["Health"])
    async def root():
        """API root endpoint."""
        return {
            "message": "Neo4j Knowledge Graph API",
            "docs": "/docs",
            "health": "/health"
        }
    
    @app.get("/health", tags=["Health"])
    async def health_check(db: Neo4jConnection = Depends(get_db)):
        """Check database health."""
        health = db.health_check()
        if health["status"] == "healthy":
            return health
        raise HTTPException(status_code=503, detail=health)
    
    # =========================================================
    # NODE CRUD ENDPOINTS
    # =========================================================
    
    @app.post("/nodes", tags=["Nodes"])
    async def create_node(
        node: NodeCreate,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Create a new node.
        
        Creates a node with specified labels and properties.
        """
        labels_str = ":".join(node.labels)
        prop_keys = list(node.properties.keys())
        prop_str = ", ".join([f"{k}: ${k}" for k in prop_keys])
        
        query = f"""
        CREATE (n:{labels_str} {{{prop_str}}})
        SET n.createdAt = datetime()
        RETURN n
        """
        
        try:
            result = db.execute_query(query, node.properties)
            return {"status": "created", "node": result[0]["n"] if result else None}
        except Exception as e:
            raise HTTPException(status_code=400, detail=str(e))
    
    @app.get("/nodes/{label}", tags=["Nodes"])
    async def get_nodes(
        label: str,
        skip: int = Query(0, ge=0),
        limit: int = Query(100, ge=1, le=1000),
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Get nodes by label with pagination.
        """
        query = f"""
        MATCH (n:{label})
        RETURN n
        ORDER BY n.id
        SKIP $skip
        LIMIT $limit
        """
        
        result = db.execute_query(query, {"skip": skip, "limit": limit})
        return {"nodes": [r["n"] for r in result], "count": len(result)}
    
    @app.get("/nodes/{label}/{node_id}", tags=["Nodes"])
    async def get_node_by_id(
        label: str,
        node_id: str,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Get a specific node by ID.
        """
        query = f"""
        MATCH (n:{label} {{id: $id}})
        RETURN n
        """
        
        result = db.execute_query(query, {"id": node_id})
        if not result:
            raise HTTPException(status_code=404, detail="Node not found")
        return {"node": result[0]["n"]}
    
    @app.put("/nodes/{label}/{node_id}", tags=["Nodes"])
    async def update_node(
        label: str,
        node_id: str,
        update: NodeUpdate,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Update node properties.
        """
        set_clauses = [f"n.{k} = ${k}" for k in update.properties.keys()]
        set_clauses.append("n.updatedAt = datetime()")
        set_str = ", ".join(set_clauses)
        
        query = f"""
        MATCH (n:{label} {{id: $node_id}})
        SET {set_str}
        RETURN n
        """
        
        params = {"node_id": node_id, **update.properties}
        result = db.execute_query(query, params)
        
        if not result:
            raise HTTPException(status_code=404, detail="Node not found")
        return {"status": "updated", "node": result[0]["n"]}
    
    @app.delete("/nodes/{label}/{node_id}", tags=["Nodes"])
    async def delete_node(
        label: str,
        node_id: str,
        detach: bool = Query(True, description="Also delete relationships"),
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Delete a node.
        """
        if detach:
            query = f"""
            MATCH (n:{label} {{id: $id}})
            DETACH DELETE n
            RETURN count(*) as deleted
            """
        else:
            query = f"""
            MATCH (n:{label} {{id: $id}})
            DELETE n
            RETURN count(*) as deleted
            """
        
        try:
            result = db.execute_query(query, {"id": node_id})
            deleted = result[0]["deleted"] if result else 0
            if deleted == 0:
                raise HTTPException(status_code=404, detail="Node not found")
            return {"status": "deleted", "count": deleted}
        except Exception as e:
            if "still has relationships" in str(e):
                raise HTTPException(
                    status_code=400,
                    detail="Node has relationships. Use detach=true to delete."
                )
            raise HTTPException(status_code=400, detail=str(e))
    
    # =========================================================
    # RELATIONSHIP ENDPOINTS
    # =========================================================
    
    @app.post("/relationships", tags=["Relationships"])
    async def create_relationship(
        rel: RelationshipCreate,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Create a relationship between two nodes.
        """
        props_str = ""
        if rel.properties:
            prop_parts = [f"{k}: $prop_{k}" for k in rel.properties.keys()]
            props_str = " {" + ", ".join(prop_parts) + "}"
        
        query = f"""
        MATCH (a:{rel.from_label} {{id: $from_id}})
        MATCH (b:{rel.to_label} {{id: $to_id}})
        CREATE (a)-[r:{rel.rel_type}{props_str}]->(b)
        SET r.createdAt = datetime()
        RETURN type(r) as type, properties(r) as properties
        """
        
        params = {
            "from_id": rel.from_id,
            "to_id": rel.to_id
        }
        if rel.properties:
            for k, v in rel.properties.items():
                params[f"prop_{k}"] = v
        
        try:
            result = db.execute_query(query, params)
            if not result:
                raise HTTPException(status_code=404, detail="One or both nodes not found")
            return {"status": "created", "relationship": result[0]}
        except Exception as e:
            raise HTTPException(status_code=400, detail=str(e))
    
    @app.get("/relationships/{node_label}/{node_id}", tags=["Relationships"])
    async def get_node_relationships(
        node_label: str,
        node_id: str,
        direction: str = Query("both", regex="^(in|out|both)$"),
        rel_type: Optional[str] = None,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Get relationships for a node.
        """
        rel_pattern = f":{rel_type}" if rel_type else ""
        
        if direction == "out":
            pattern = f"-[r{rel_pattern}]->"
        elif direction == "in":
            pattern = f"<-[r{rel_pattern}]-"
        else:
            pattern = f"-[r{rel_pattern}]-"
        
        query = f"""
        MATCH (n:{node_label} {{id: $id}}){pattern}(other)
        RETURN 
            type(r) as type,
            properties(r) as properties,
            labels(other)[0] as otherLabel,
            other.id as otherId,
            other.name as otherName
        """
        
        result = db.execute_query(query, {"id": node_id})
        return {"relationships": result}
    
    # =========================================================
    # QUERY ENDPOINTS
    # =========================================================
    
    @app.post("/query", tags=["Query"])
    async def execute_query(
        query: CypherQuery,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Execute a Cypher query.
        
        WARNING: Only for trusted users. Validate queries in production.
        """
        # Basic safety checks
        query_lower = query.query.lower()
        if any(dangerous in query_lower for dangerous in ["drop", "delete", "remove", "detach"]):
            if "where" not in query_lower:
                raise HTTPException(
                    status_code=400,
                    detail="Destructive operations require WHERE clause"
                )
        
        try:
            result = db.execute_query(query.query, query.parameters or {})
            return {"results": result, "count": len(result)}
        except Exception as e:
            raise HTTPException(status_code=400, detail=str(e))
    
    @app.post("/search", tags=["Query"])
    async def search_nodes(
        search: SearchQuery,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Search nodes by text.
        """
        label_filter = ""
        if search.labels:
            label_filter = ":" + ":".join(search.labels)
        
        query = f"""
        MATCH (n{label_filter})
        WHERE toLower(n.name) CONTAINS toLower($query)
           OR toLower(n.id) CONTAINS toLower($query)
        RETURN n
        ORDER BY n.name
        LIMIT $limit
        """
        
        result = db.execute_query(
            query,
            {"query": search.query, "limit": search.limit}
        )
        return {"results": [r["n"] for r in result], "count": len(result)}
    
    # =========================================================
    # VECTOR SEARCH ENDPOINTS
    # =========================================================
    
    @app.post("/vector-search", tags=["Vector Search"])
    async def vector_similarity_search(
        search: VectorSearchQuery,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Perform vector similarity search.
        
        Requires:
        - Vector index to be created
        - Embedding model available
        """
        # In production, you would embed the text here
        # For now, return a helpful message
        return {
            "message": "Vector search requires embedding the query text first",
            "index_name": search.index_name,
            "query_text": search.text,
            "top_k": search.top_k,
            "hint": "Use the GraphRAG class for full vector search functionality"
        }
    
    # =========================================================
    # RAG ENDPOINTS
    # =========================================================
    
    @app.post("/rag/query", tags=["RAG"])
    async def rag_query(
        query: RAGQuery,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Execute a Graph RAG query.
        
        Combines vector search with graph traversal to answer questions.
        """
        # This would use the GraphRAG class in production
        return {
            "question": query.question,
            "answer": "[RAG system would generate answer here]",
            "context": {
                "chunks_retrieved": query.top_k,
                "graph_context_included": query.include_graph_context
            },
            "note": "Connect GraphRAG class for full functionality"
        }
    
    # =========================================================
    # BATCH ENDPOINTS
    # =========================================================
    
    @app.post("/batch/nodes", tags=["Batch"])
    async def batch_create_nodes(
        batch: BatchNodeCreate,
        background_tasks: BackgroundTasks,
        db: Neo4jConnection = Depends(get_db)
    ):
        """
        Create multiple nodes in a batch.
        """
        created = 0
        errors = []
        
        for node in batch.nodes:
            try:
                labels_str = ":".join(node.labels)
                prop_keys = list(node.properties.keys())
                prop_str = ", ".join([f"{k}: ${k}" for k in prop_keys])
                
                query = f"""
                CREATE (n:{labels_str} {{{prop_str}}})
                SET n.createdAt = datetime()
                RETURN n.id as id
                """
                
                db.execute_query(query, node.properties)
                created += 1
            except Exception as e:
                errors.append({"node": node.properties.get("id"), "error": str(e)})
        
        return {
            "status": "completed",
            "created": created,
            "errors": errors,
            "total": len(batch.nodes)
        }
    
    # =========================================================
    # STATISTICS ENDPOINTS
    # =========================================================
    
    @app.get("/stats", tags=["Statistics"])
    async def get_database_stats(db: Neo4jConnection = Depends(get_db)):
        """
        Get database statistics.
        """
        query = """
        CALL apoc.meta.stats() YIELD nodeCount, relCount, labelCount, relTypeCount
        RETURN nodeCount, relCount, labelCount, relTypeCount
        """
        
        try:
            result = db.execute_query(query)
            return result[0] if result else {"error": "Could not get stats"}
        except:
            # Fallback without APOC
            fallback_query = """
            MATCH (n)
            WITH count(n) as nodeCount
            MATCH ()-[r]->()
            WITH nodeCount, count(r) as relCount
            RETURN nodeCount, relCount
            """
            result = db.execute_query(fallback_query)
            return result[0] if result else {"nodeCount": 0, "relCount": 0}
    
    @app.get("/stats/labels", tags=["Statistics"])
    async def get_label_stats(db: Neo4jConnection = Depends(get_db)):
        """
        Get node counts by label.
        """
        query = """
        MATCH (n)
        RETURN labels(n)[0] as label, count(n) as count
        ORDER BY count DESC
        """
        
        result = db.execute_query(query)
        return {"labels": result}


# ============================================================
# SERVER STARTUP
# ============================================================

def run_server(host: str = "0.0.0.0", port: int = 8000):
    """Run the API server."""
    if not FASTAPI_AVAILABLE:
        print("FastAPI not installed. Run: pip install fastapi uvicorn")
        return
    
    print(f"""
    ╔══════════════════════════════════════════════════════════════╗
    ║           Neo4j Knowledge Graph API Server                   ║
    ╠══════════════════════════════════════════════════════════════╣
    ║                                                              ║
    ║  Server: http://{host}:{port}                                  ║
    ║  Docs:   http://{host}:{port}/docs                             ║
    ║  ReDoc:  http://{host}:{port}/redoc                            ║
    ║                                                              ║
    ╚══════════════════════════════════════════════════════════════╝
    """)
    
    uvicorn.run(app, host=host, port=port)


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Neo4j REST API Server")
    print("=" * 60)
    
    if FASTAPI_AVAILABLE:
        print("""
    Available Endpoints:
    
    HEALTH:
      GET  /health              - Database health check
    
    NODES:
      POST   /nodes             - Create a node
      GET    /nodes/{label}     - Get nodes by label
      GET    /nodes/{label}/{id} - Get specific node
      PUT    /nodes/{label}/{id} - Update node
      DELETE /nodes/{label}/{id} - Delete node
    
    RELATIONSHIPS:
      POST /relationships       - Create relationship
      GET  /relationships/{label}/{id} - Get node relationships
    
    QUERY:
      POST /query              - Execute Cypher query
      POST /search             - Search nodes by text
      POST /vector-search      - Vector similarity search
    
    RAG:
      POST /rag/query          - Execute RAG query
    
    BATCH:
      POST /batch/nodes        - Batch create nodes
    
    STATISTICS:
      GET  /stats              - Database statistics
      GET  /stats/labels       - Node counts by label
        """)
        
        print("\n   Starting server...")
        run_server()
    else:
        print("   FastAPI not available. Install with:")
        print("   pip install fastapi uvicorn")
